import asyncio
import os
import traceback
from typing import List

from fastapi import APIRouter, Depends, UploadFile, Form, BackgroundTasks

from app.api.response import Response, response_docs
from app.business.milvus_content import MilvusContentCreateBiz, MilvusContentExistsBiz, MilvusContentDeleteBiz, \
    MilvusContentListBiz
from app.business.milvus_index import (
    MilvusIndexDetailBiz, MilvusIndexCreateBiz, MilvusIndexListBiz, MilvusIndexDeleteBiz, MilvusIndexUpdateBiz,
)
from app.api.status import Status
from app.datatype.milvus_content import MilvusContentDeleteMdl, MilvusFilePath, MilvusContentCreateMdl, \
    MilvusContentCreateBatchMdl
from app.initializer import g
from app.middleware.auth import JWTUser, get_current_user
from app.utils.milvus_do import MilvusDo

milvus_index_router = APIRouter()


#
# @milvus_index_router.get(
#     path="/milvus_index/{milvus_index_id}",
#     summary="milvus_indexDetail",
#     responses=response_docs(
#         model=MilvusIndexDetailBiz,
#     ),
# )
# async def detail(
#         milvus_index_id: str,
#         current_user: JWTUser = Depends(get_current_user),
# ):
#     try:
#         milvus_index_biz = MilvusIndexDetailBiz(id=milvus_index_id)
#         data = await milvus_index_biz.detail()
#         if not data:
#             return Response.failure(msg="未匹配到记录", status=Status.RECORD_NOT_EXIST_ERROR)
#     except Exception as e:
#         g.logger.error(traceback.format_exc())
#         return Response.failure(msg="milvus_indexDetail失败", error=e)
#     return Response.success(data=data)

# 创建知识库
@milvus_index_router.post(
    path="/createMilvusIndex",
    summary="创建知识库V2",
    responses=response_docs(data={
        "id": int,
    }),
)
async def create(
        milvus_biz: MilvusIndexCreateBiz,
        current_user: JWTUser = Depends(get_current_user),
):
    try:
        user_id = current_user.id
        milvus_biz.create_id = user_id
        session_id = await milvus_biz.create()
        if not session_id:
            return Response.failure(msg="创建失败,知识库名重复", status=Status.RECORD_EXISTS_ERROR)
    except Exception as e:
        g.logger.error(traceback.format_exc())
        return Response.failure(msg="创建失败", error=e)
    return Response.success(data={"id": session_id})


from app.datatype.milvus_index import MilvusIndexListMdl

@milvus_index_router.post(
    path="/createMilvusUserIndex",
    summary="创建个人知识库V2",
    responses=response_docs(data={
        "id": int,
    }),
)
async def create(
        milvus_biz: MilvusIndexCreateBiz,
        current_user: JWTUser = Depends(get_current_user),
):
    try:
        user_id = current_user.id
        milvus_biz.create_id = user_id
        milvus_biz.user_id = user_id
        session_id = await milvus_biz.create()
        if not session_id:
            return Response.failure(msg="创建失败,知识库名重复", status=Status.RECORD_EXISTS_ERROR)
    except Exception as e:
        g.logger.error(traceback.format_exc())
        return Response.failure(msg="创建失败", error=e)
    return Response.success(data={"id": session_id})


from app.datatype.milvus_index import MilvusIndexListMdl
# 获取知识库列表
@milvus_index_router.post(
    path="/milvusIndexList",
    summary="获取知识库列表管理端V2",
    responses=response_docs(
        model=MilvusIndexListMdl,
        is_listwrap=True,
    ),
)
async def lst(
        milvus_biz: MilvusIndexListBiz,
        current_user: JWTUser = Depends(get_current_user),  # 认证
):
    try:
        data = await milvus_biz.lst(current_user =current_user)
        if not data:
            return Response.failure(msg="未匹配到记录", status=Status.RECORD_NOT_EXIST_ERROR)
    except Exception as e:
        g.logger.error(traceback.format_exc())
        return Response.failure(msg="获取知识库列表失败", error=e)
    return Response.success(data=data)
# 获取知识库列表
@milvus_index_router.post(
    path="/milvusIndexListPersonage",
    summary="获取个人的知识库列表V2",
    responses=response_docs(
        model=MilvusIndexListMdl,
        is_listwrap=True,
    ),
)
async def lstP(
        milvus_biz: MilvusIndexListBiz,
        current_user: JWTUser = Depends(get_current_user),  # 认证
):
    try:
        data = await milvus_biz.lstPer(current_user =current_user)
        if not data:
            return Response.failure(msg="未匹配到记录", status=Status.RECORD_NOT_EXIST_ERROR)
    except Exception as e:
        g.logger.error(traceback.format_exc())
        return Response.failure(msg="获取知识库列表失败", error=e)
    return Response.success(data=data)

# 删除知识库
@milvus_index_router.delete(
    path="/deleteMilvusIndex",
    summary="删除知识库",
    responses=response_docs(data={
        "status": str,
    }),
)
async def delete(
        milvus_biz: MilvusIndexDeleteBiz,
        current_user: JWTUser = Depends(get_current_user),
):
    try:
        status = await milvus_biz.delete()
        if not status:
            return Response.failure(msg="删除失败", status=Status.RECORD_NOT_EXIST_ERROR)
    except Exception as e:
        g.logger.error(traceback.format_exc())
        return Response.failure(msg="删除知识库失败", error=e)
    return Response.success(data={"status": "success"})


# 修改知识库
@milvus_index_router.put(
    path="/updateMilvusIndex",
    summary="修改知识库",
    responses=response_docs(data={
        "status": int,
    }),
)
async def update(
        milvus_biz: MilvusIndexUpdateBiz,
        current_user: JWTUser = Depends(get_current_user),
):
    try:
        status = await milvus_biz.update()
        if not status:
            return Response.failure(msg="更新失败", status=Status.RECORD_NOT_EXIST_ERROR)
    except Exception as e:
        g.logger.error(traceback.format_exc())
        return Response.failure(msg="更新知识库失败", error=e)
    return Response.success(data={"status": status})


@milvus_index_router.post("/upload/fileToMilvus",
                          summary="文件单个上传到Milvus",
                          responses=response_docs(data={
                              "status": str,
                          }),
                          )
async def upload_file(
        milvus_biz: MilvusContentCreateBiz,
        current_user: JWTUser = Depends(get_current_user),
):
    async def process_upload(milvus_biz):
        file_path = milvus_biz.file_path
        file_name = os.path.basename(file_path)
        file_size_kb = os.path.getsize(file_path) / 1024
        milvus_biz.file_size = file_size_kb
        milvus_content_id = await milvus_biz.create()
        milvus_instance = MilvusDo()
        try:
            status =milvus_instance.upload_file_to_milvus(milvus_biz.milvus_index, file_path, file_name)
        except Exception as e:

            status = "failed"
        if status=="failed":
            await milvus_biz.update_upload_status_fail(milvus_content_id=milvus_content_id)
        else:
            await milvus_biz.update_upload_status(milvus_content_id=milvus_content_id)
    asyncio.create_task(process_upload(milvus_biz))
    return Response.success(data={"status": "success"})


# 文件上传到指定知识库
@milvus_index_router.post("/upload/uploadFileCollection",
                          summary="文件上传到指定知识库获取文件path",
                          responses=response_docs(data={
                              "filePath": str,
                          }),
                          )
async def upload_file(
        file_upload: UploadFile,
        milvus_index: str = Form(..., description="索引名称"),
        current_user: JWTUser = Depends(get_current_user),
):
    # 目标文件夹路径
    folder_path = "collection/" + milvus_index
    file_name = file_upload.filename
    # 检查文件是否存在知识库
    milvus_find_count = MilvusContentExistsBiz()
    count = await milvus_find_count.exists(milvus_index=milvus_index, file_name=file_name)
    if count > 0:
        return Response.failure(msg="文件已存在,请重新修改文件名", status=Status.RECORD_EXISTS_ERROR)
    # 检查文件夹是否存在，如果不存在则创建
    if not os.path.exists(folder_path):
        os.makedirs(folder_path)
    a = await  file_upload.read()
    # 存储到当前文件夹中的特定文件夹

    save_path = folder_path + "/" + file_name
    with open(save_path, 'wb') as file:
        file.write(a)

    #返回文件路径

    return Response.success(data={"filePath":save_path})

#删除知识库文件
@milvus_index_router.delete("/delete/deleteFileCollection",
                            summary="删除知识库文件",
                            responses=response_docs(data={
                                "status": str,
                            }),
                            )
async def delete_file(
        milvus_delete: MilvusContentDeleteBiz,
        current_user: JWTUser = Depends(get_current_user),
):
    try:
        status = await milvus_delete.delete()
        if not status:
            return Response.failure(msg="删除失败,文件正在上传中", status=Status.RECORD_NOT_EXIST_ERROR)
    except Exception as e:
        g.logger.error(traceback.format_exc())
        return Response.failure(msg="删除知识库文件失败", error=e)
    return Response.success(data={"status": "success"})


#获取知识库文件列表
@milvus_index_router.post(
    path="/milvusContentList",
    summary="获取知识库文件列表",
    responses=response_docs(
        model=MilvusContentListBiz,
        is_listwrap=True,
    ),
)
async def milvus_content_list(
        milvus_biz: MilvusContentListBiz,
        current_user: JWTUser = Depends(get_current_user),
):
    try:
        data = await milvus_biz.lst()
        if not data:
            return Response.failure(msg="未匹配到记录", status=Status.RECORD_NOT_EXIST_ERROR)
    except Exception as e:
        g.logger.error(traceback.format_exc())
        return Response.failure(msg="获取知识库文件列表失败", error=e)
    return Response.success(data=data)

#通过path 获取文档 预览文档
@milvus_index_router.post("/downloadFileByPath",
                         summary="通过filePath下载文件"
                         )
async def download_file_by_path(
        mil_path: MilvusFilePath
):
    file_path  =mil_path.path
    if not os.path.exists(file_path):
        return Response.failure(msg="文件不存在", status=Status.RECORD_NOT_EXIST_ERROR)
    file_name = os.path.basename(file_path)
    from fastapi.responses import FileResponse
    return FileResponse(path=file_path, filename=file_name, media_type='application/octet-stream')

@milvus_index_router.post("/upload/fileToMilvusBatch",
                          summary="文件多文件上传到Milvus",
                          responses=response_docs(data={
                              "status": str,
                          }),
                          )
async def upload_file(
        milvus_biz_list: List[MilvusContentCreateBatchMdl],
        current_user: JWTUser = Depends(get_current_user),
):

    async def process_upload(milvus_biz_list):
        for milvus_biz in milvus_biz_list:
            if not milvus_biz.file_path or not os.path.exists(milvus_biz.file_path):
                continue

            file_path = milvus_biz.file_path
            file_name = os.path.basename(file_path)
            file_size_kb = os.path.getsize(file_path) / 1024

            # 入库
            milvus_biz_create = MilvusContentCreateBiz()
            milvus_biz_create.file_size = file_size_kb
            milvus_biz_create.milvus_index = milvus_biz.milvus_index
            milvus_biz_create.milvus_index_id = milvus_biz.milvus_index_id
            milvus_biz_create.file_name = file_name
            milvus_biz_create.file_path = file_path
            milvus_content_id = await milvus_biz_create.create()

            milvus_biz.milvus_content_id = milvus_content_id
        milvus_create_biz = MilvusContentCreateBiz()
        for milvus_biz in milvus_biz_list:
            file_path = milvus_biz.file_path
            file_name = os.path.basename(file_path)
            try:
                milvus_instance = MilvusDo()
                status =milvus_instance.upload_file_to_milvus(milvus_biz.milvus_index, file_path, file_name)
            except Exception as e:
                status = "failed"

            if status=="failed":
                print(f"Failed to upload milvus_content_id  {milvus_biz.milvus_content_id} to {milvus_biz.milvus_index}")
                await milvus_create_biz.update_upload_status_fail(milvus_content_id=milvus_biz.milvus_content_id)
            else:
                print(f"Successfully uploaded milvus_content_id {milvus_biz.milvus_content_id} to {milvus_biz.milvus_index}")
                await milvus_create_biz.update_upload_status(milvus_content_id=milvus_biz.milvus_content_id)
    asyncio.create_task(process_upload(milvus_biz_list))
    return Response.success(data={"status": "success"})
